
import os

from langchain.agents import create_react_agent, AgentExecutor, Tool
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables.history import RunnableWithMessageHistory

os.environ["SERPAPI_API_KEY"] = '0fA485b7947AA6becd79b9098c191073'

llm = Tongyi()


def buy_goods(mes):
    print(mes)
    return "购买成功"


def temperature_weather(mes):
    print(mes)
    return "今天天气很不错，温度25度"


def get_emotion(mes):
    print(mes)
    return "很高兴"


goods = Tool(func=buy_goods, name="购物", description="购买商品")
weather = Tool(func=temperature_weather, name="天气", description="查询天气")
emotion = Tool(func=get_emotion, name="情感", description="获取情感")
tools = [goods, weather, emotion]

prompt_template = """
您是一个助手，可以访问以下工具来帮助解决问题：
{tools}
请按照以下格式回答问题：

Question: {input}
Thought: 如果问题是关于天气或温度，我应该使用temperature_weather工具；如果问题是关于情感，我应该使用get_emotion工具;如果问题是关于商品购物的，我应该使用buy_goods工具。
Action: 使用工具的名称 ({tool_names})
Action Input:{input}
Thought: 输出工具得返回结果
Observation: 
Final Answer: 最终答案
Final Answer: the final answer to the original input question

Begin!

Question: {input}
Thought:{agent_scratchpad}
现在，请根据问题开始思考并执行相应的动作。

"""
prompt = ChatPromptTemplate.from_template(prompt_template)

agent = create_react_agent(llm, tools, prompt)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

message_history = ChatMessageHistory()

agent_with_chat_history = RunnableWithMessageHistory(
    agent_executor,
    lambda session_id: message_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)

res = agent_with_chat_history.invoke(
    {"input": "我要买手机"},
    config={"configurable": {"session_id": "session-10086"}},
)

print(res)

res = agent_with_chat_history.invoke(
    {"input": "天气怎么样？"},
    config={"configurable": {"session_id": "session-10086"}},
)
print("-")
print(res)

res = agent_with_chat_history.invoke(
    {"input": "你怎么样？"},
    config={"configurable": {"session_id": "session-10086"}},
)
print("-")
print(res)
